
class Graph:
    '''
    计算图类
    '''

    def __init__(self):
        self.nodes = [] # 计算图内的节点的列表
        self.name_scope = None

    def add_node(self, node):
        '''
        添加节点
        :param node:
        :return:
        '''
        self.nodes.append(node)

    def clear_jacobi(self):
        '''
        清空计算图中全部节点的雅可比矩阵
        :return:
        '''
        for node in self.nodes:
            node.clear_jacobi()

    def reset_value(self):
        '''
        重置计算图中全部节点的值
        :return:
        '''
        for node in self.nodes:
            node.reset_value(False)

    def node_count(self):
        return len(self.nodes)

    def draw(self, ax=None):
        try:
            import networkx as nx
            import matplotlib.pyplot as plt
            from matplotlib.colors import ListedColormap
            import numpy as np
        except:
            raise Exception("Need Module networkx")

        G = nx.Graph()

        already = []
        labels = {}
        for node in self.nodes:
            G.add_node(node)
            labels[node] = node.__class__.__name__ + ("({:s})".format(str(node.dim)) if hasattr(node, "dim") else "")\
                           + ("\n[{:.3f}]".format(np.linalg.norm(node.jacobi)) if node.jacobi is not None else "")

            for c in node.get_children():
                if {node, c} not in already:
                    G.add_edge(node, c)
                    already.append({node, c})
            for p in node.get_parents():
                if {node, p} not in already:
                    G.add_edge(node, p)
                    already.append({node, p})

        if ax is None:
            fig = plt.figure(figsize=(12, 12))
            ax = fig.add_subplot(111)

        ax.clear()
        ax.axis("on")
        ax.grid(True)

        pos = nx.sprint_layout(G, seed=42)

        '''
        未完待续，这个draw函数目前还用不到
        '''

default_graph = Graph()